module mxu_mat_array (
    input logic                                 clk,
    input logic                                 rst_n,
          mxu_mat_ctrl_mat_array_if.mat_ctrl_in u_mxu_mat_ctrl_mat_array_if,
          mxu_psb_dat_if.out                    mat_psb_dat
);

    logic mxu_mat_array_en;
    assign mxu_mat_array_en = mat_psb_dat.rdy;

    genvar m, n;
    pack_16B_t       a_delayed_in   [`M0-1:0];
    pack_16B_t       b_delayed_in   [`N0-1:0];
    fp32_t           c_delayed_in   [`M0-1:0] [`N0-1:0];
    logic      [2:0] ctrl_delayed_in[`M0-1:0] [`N0-1:0];

    // 修正类型：bf16_t → pack_16B_t
    // 扩展 a_net 和 b_net 以容纳边界 PE 的输出
    pack_16B_t       a_net          [`M0-1:0] [  `N0:0];  // 增加一列 (n=N0) 用于边界输出
    pack_16B_t       b_net          [  `M0:0] [`N0-1:0];  // 增加一行 (m=M0) 用于边界输出
    fp32_t           d_out          [`M0-1:0] [`N0-1:0];
    logic            out_vld        [`M0-1:0] [`N0-1:0];  // 新增：接收 valid_out

    fp32_t           d_delayed_out  [`M0-1:0] [`N0-1:0];
    logic            fsum_vld;

    generate
        for (m = 0; m < `M0; m = m + 1) begin : gen_m_delay
            common_delay_line #(
                .WIDTH($bits(u_mxu_mat_ctrl_mat_array_if.a[m])),  //修正信号名
                .DEPTH(m)
            ) u_a_delay_line (
                .clk     (clk),
                .rst_n   (rst_n),
                .en      (mxu_mat_array_en),
                .data_in (u_mxu_mat_ctrl_mat_array_if.a[m]),
                .data_out(a_delayed_in[m])
            );
        end
    endgenerate
    generate
        for (n = 0; n < `N0; n = n + 1) begin : gen_n_delay
            common_delay_line #(
                .WIDTH($bits(u_mxu_mat_ctrl_mat_array_if.b[n])),
                .DEPTH(n)
            ) u_b_delay_line (
                .clk     (clk),
                .rst_n   (rst_n),
                .en      (mxu_mat_array_en),
                .data_in (u_mxu_mat_ctrl_mat_array_if.b[n]),
                .data_out(b_delayed_in[n])
            );
        end
    endgenerate
    generate
        for (m = 0; m < `M0; m = m + 1) begin : gen_cd_delay_m
            for (n = 0; n < `N0; n = n + 1) begin : gen_cd_delay_n
                common_delay_line #(
                    .WIDTH($bits(fp32_t)),
                    .DEPTH(m + n)
                ) u_c_delay_line (
                    .clk     (clk),
                    .rst_n   (rst_n),
                    .en      (mxu_mat_array_en),
                    .data_in (u_mxu_mat_ctrl_mat_array_if.c[m][n]),
                    .data_out(c_delayed_in[m][n])
                );
                common_delay_line #(
                    .WIDTH(3),
                    .DEPTH(m + n)  //ctrl_delay用于传输mode，到底应该是a*b还是a*b+c还是a*b+psum
                ) u_ctrl_delay_line (
                    .clk     (clk),
                    .rst_n   (rst_n),
                    .en      (mxu_mat_array_en),
                    .data_in ({u_mxu_mat_ctrl_mat_array_if.mode, u_mxu_mat_ctrl_mat_array_if.vld}),
                    .data_out(ctrl_delayed_in[m][n])
                );
                common_delay_line #(
                    .WIDTH($bits(fp32_t)),
                    .DEPTH(`M0 + `N0 - m - n)
                ) u_d_delay_line (
                    .clk     (clk),
                    .rst_n   (rst_n),
                    .en      (mxu_mat_array_en),
                    .data_in (d_out[m][n]),
                    .data_out(d_delayed_out[m][n])
                );
            end
        end
    endgenerate

    // last_k的数据出现在输出端口的delay是固定的，因为从输入数据到可以输出数据，中间需要经过的延迟链就是输入延迟链和计算延迟链，理论上是`M0 + `N0 + `MAT_LATENCY，但是不是很准确，待调
    common_delay_line #(
        .DATA_WIDTH  (1),
        .DELAY_CYCLES(`M0 + `N0 + `MAT_LATENCY)
    ) u_last_k_delay_line (
        .clk  (clk),
        .rst_n(rst_n),
        .en   (mxu_mat_array_en),
        .din  (u_mxu_mat_ctrl_mat_array_if.dat_vld && u_mxu_mat_ctrl_mat_array_if.last_k),
        .dout (fsum_vld)
    );

    logic exception_array                    [`M0-1:0][`N0-1:0];  // 存储每个 PE 的异常
    logic exception_any;  // 任一 PE 异常

    generate
        for (m = 0; m < `M0; m = m + 1) begin : gen_a_net_assign
            assign a_net[m][0] = a_delayed_in[m];
        end
        for (n = 0; n < `N0; n = n + 1) begin : gen_b_net_assign
            assign b_net[0][n] = b_delayed_in[n];
        end
    endgenerate

    generate
        for (m = 0; m < `M0; m = m + 1) begin : gen_m
            for (n = 0; n < `N0; n = n + 1) begin : gen_n
                mxu_mat u_mxu_mat (
                    .clk   (clk),
                    .rst_n (rst_n),
                    .a_in  (a_net[m][n]),
                    .b_in  (b_net[m][n]),
                    .c_in  (c_delayed_in[m][n]),
                    .in_vld(ctrl_delayed_in[m][n][0]),
                    .mode  (ctrl_delayed_in[m][n][2:1]),  // 00: a*b; 01: a*b+c; 10: a*b+psum
                    .a_out (a_net[m][n+1]),
                    .b_out (b_net[m+1][n]),
                    .d_out (d_out[m][n]),
                    .out_vld(out_vld[m][n]),
                    .exception(exception_array[m][n])    // 收集异常
                );
            end
        end
    endgenerate

    // 汇总异常 - 手动展开二维数组
    logic [`M0*`N0-1:0] exception_flat;
    generate
        for (m = 0; m < `M0; m = m + 1) begin : gen_exception_m
            for (n = 0; n < `N0; n = n + 1) begin : gen_exception_n
                assign exception_flat[m*`N0+n] = exception_array[m][n];
            end
        end
    endgenerate
    assign exception_any   = |exception_flat;

    assign mat_psb_dat.vld = fsum_vld;
    assign mat_psb_dat.dat = d_delayed_out;

endmodule
